import numpy as np
import networkx as nx

from networkx.generators.community import LFR_benchmark_graph
import pickle
import os

import cdlib



if __name__ == "__main__":
    n_graphs = 10
    tries=50
    # mut = np.arange(0.25,0.6,0.05)
    mut = np.arange(0.6,0.75,0.05)
    
    if os.path.isfile('LFR_benchmark.pkl'):
        res = pickle.load(open('LFR_benchmark.pkl','rb'))
    else:
        res = {'graphs': [],
               'mut': [],
               'n': [],
               'ground_truth': []}
    
    mut = np.tile(mut,n_graphs)
    
    #main loop
    for mut_ in mut:
    
        #generate graph
        count = 0
        for t in range(tries):
            count += 1
            print('Try ' + str(count))
            if count == tries:
                raise Exception("Graph generation failed.")
            try:
                G = LFR_benchmark_graph(n=500, 
                                tau1=3, 
                                tau2=1.5, 
                                mu=mut_, 
                                average_degree=20,
                                max_degree=50,
                                min_community=10, 
                                max_community=50)
                
                print('Graph generation ok.')
                
                largest_cc = max(nx.connected_components(G), key=len)
                G = G.subgraph(largest_cc)
                G = nx.convert_node_labels_to_integers(G)
        
                C = {frozenset(G.nodes[v]["community"]) for v in G}
                C = [list(x) for x in C]
        
                ground_truth = cdlib.NodeClustering(communities=C, graph=G, method_name="reference")
        
                #collect metadata
                res['graphs'].append(G)
                res['mut'].append(mut_)
                res['ground_truth'].append(ground_truth)
                
                break
            except:
                continue
        
    #saving
    pickle.dump(res, open('LFR_benchmark.pkl','wb'))